#include "phasar/PhasarLLVM/DataFlow/IfdsIde/Problems/IFDSTaintAnalysis.h"

#include "phasar/DataFlow/IfdsIde/Solver/IFDSSolver.h"
#include "phasar/PhasarLLVM/ControlFlow/LLVMBasedICFG.h"
#include "phasar/PhasarLLVM/HelperAnalyses.h"
#include "phasar/PhasarLLVM/Pointer/LLVMAliasSet.h"
#include "phasar/PhasarLLVM/SimpleAnalysisConstructor.h"
#include "phasar/PhasarLLVM/TaintConfig/LLVMTaintConfig.h"
#include "phasar/PhasarLLVM/TaintConfig/TaintConfigBase.h"
#include "phasar/Utils/DebugOutput.h"

#include "llvm/ADT/StringRef.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Value.h"

#include "SrcCodeLocationEntry.h"
#include "TestConfig.h"
#include "gtest/gtest.h"

using namespace psr;
using namespace psr::unittest;

/* ============== TEST FIXTURE ============== */

class IFDSTaintAnalysisTest : public ::testing::Test {
protected:
  static constexpr auto PathToLlFiles =
      PHASAR_BUILD_SUBFOLDER("taint_analysis/");
  static inline const std::vector<std::string> EntryPoints = {"main"};

  std::optional<HelperAnalyses> HA;

  std::optional<IFDSTaintAnalysis> TaintProblem;
  std::optional<LLVMTaintConfig> TSF;

  static LLVMTaintConfig getDefaultConfig() {
    auto SourceCB = [](const llvm::Instruction *Inst) {
      std::set<const llvm::Value *> Ret;
      if (const auto *Call = llvm::dyn_cast<llvm::CallBase>(Inst);
          Call && Call->getCalledFunction() &&
          Call->getCalledFunction()->getName() == "_Z6sourcev") {
        Ret.insert(Call);
      }
      return Ret;
    };
    auto SinkCB = [](const llvm::Instruction *Inst) {
      std::set<const llvm::Value *> Ret;
      if (const auto *Call = llvm::dyn_cast<llvm::CallBase>(Inst);
          Call && Call->getCalledFunction() &&
          Call->getCalledFunction()->getName() == "_Z4sinki") {
        assert(Call->arg_size() > 0);
        Ret.insert(Call->getArgOperand(0));
      }
      return Ret;
    };
    return LLVMTaintConfig(std::move(SourceCB), std::move(SinkCB));
  }

  static LLVMTaintConfig getDoubleFreeConfig() {
    auto SourceCB = [](const llvm::Instruction *Inst) {
      std::set<const llvm::Value *> Ret;
      if (const auto *Call = llvm::dyn_cast<llvm::CallBase>(Inst);
          Call && Call->getCalledFunction() &&
          Call->getCalledFunction()->getName() == "free") {
        Ret.insert(Call->getArgOperand(0));
      }
      return Ret;
    };

    return LLVMTaintConfig(SourceCB, SourceCB);
  }

  void initialize(const llvm::Twine &IRFile) {
    HA.emplace(IRFile, EntryPoints);

    if (!TSF) {
      TSF = getDefaultConfig();
    }

    TaintProblem =
        createAnalysisProblem<IFDSTaintAnalysis>(*HA, &*TSF, EntryPoints);
  }

  void initialize(const llvm::Twine &IRFile, const LLVMTaintConfig *Config) {
    HA.emplace(IRFile, EntryPoints);
    TaintProblem =
        createAnalysisProblem<IFDSTaintAnalysis>(*HA, Config, EntryPoints);
  }

  using GroundTruthTy =
      std::map<TestingSrcLocation, std::set<TestingSrcLocation>>;

  template <typename LeaksTy>
  void compare(const LeaksTy &Leaks, const GroundTruthTy &GroundTruth) {
    auto GroundTruthEntries =
        convertTestingLocationSetMapInIR(GroundTruth, HA->getProjectIRDB());

    EXPECT_EQ(Leaks, GroundTruthEntries)
        << "Taint Leaks do not match:\n  Expected: "
        << PrettyPrinter{GroundTruthEntries}
        << "\n  Got: " << PrettyPrinter{Leaks};
  }

}; // Test Fixture

TEST_F(IFDSTaintAnalysisTest, TaintTest_01) {
  initialize({PathToLlFiles + "dummy_source_sink/taint_01_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{6, 3, "main"};
  auto EntryTwo = LineColFun{6, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_01_m2r) {
  initialize({PathToLlFiles + "dummy_source_sink/taint_01_cpp_m2r_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{6, 3, "main"};
  auto EntryTwo = LineColFun{5, 11, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_02) {
  initialize({PathToLlFiles + "dummy_source_sink/taint_02_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{5, 3, "main"};
  auto EntryTwo = LineColFun{5, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_03) {
  initialize({PathToLlFiles + "dummy_source_sink/taint_03_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{6, 3, "main"};
  auto EntryTwo = LineColFun{6, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_04) {
  initialize({PathToLlFiles + "dummy_source_sink/taint_04_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{6, 3, "main"};
  auto EntryTwo = LineColFun{6, 8, "main"};
  auto EntryThree = LineColFun{8, 3, "main"};
  auto EntryFour = LineColFun{8, 8, "main"};
  GroundTruthTy GroundTruth{
      {Entry, {EntryTwo}},
      {EntryThree, {EntryFour}},
  };

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_05) {
  initialize({PathToLlFiles + "dummy_source_sink/taint_05_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{6, 3, "main"};
  auto EntryTwo = LineColFun{6, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_06) {
  initialize({PathToLlFiles + "dummy_source_sink/taint_06_cpp_m2r_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{5, 3, "main"};
  auto Main0 = ArgInFun{0, "main"};
  GroundTruthTy GroundTruth{{Entry, {Main0}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_01) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_01_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{12, 3, "main"};
  auto EntryTwo = LineColFun{12, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_01_m2r) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_01_cpp_m2r_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{12, 3, "main"};
  auto EntryTwo = LineColFun{10, 14, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_02) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_02_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{11, 3, "main"};
  auto EntryTwo = LineColFun{11, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_03) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_03_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{11, 3, "main"};
  auto EntryTwo = LineColFun{11, 8, "main"};
  auto EntryThree = LineColFun{14, 3, "main"};
  auto EntryFour = LineColFun{14, 8, "main"};
  GroundTruthTy GroundTruth{
      {Entry, {EntryTwo}},
      {EntryThree, {EntryFour}},
  };

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_04) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_04_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{16, 3, "main"};
  auto EntryTwo = LineColFun{16, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_05) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_05_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{16, 3, "main"};
  auto EntryTwo = LineColFun{16, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_06) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_06_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{13, 5, "main"};
  auto EntryTwo = LineColFun{13, 10, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_07) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_07_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{14, 5, "main"};
  auto EntryTwo = LineColFun{14, 10, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_08) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_08_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{19, 3, "main"};
  auto EntryTwo = LineColFun{19, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_09) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_09_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{20, 3, "main"};
  auto EntryTwo = LineColFun{20, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_10) {
  initialize(
      {PathToLlFiles + "dummy_source_sink/taint_exception_10_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{19, 5, "main"};
  auto EntryTwo = LineColFun{19, 10, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_DoubleFree_01) {
  auto DoubleFreeConf = getDoubleFreeConfig();
  initialize({PathToLlFiles + "double_free_01_c_dbg.ll"}, &DoubleFreeConf);
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{6, 3, "main"};
  auto EntryTwo = LineColFun{6, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_DoubleFree_02) {
  auto DoubleFreeConf = getDoubleFreeConfig();
  initialize({PathToLlFiles + "double_free_02_c_dbg.ll"}, &DoubleFreeConf);
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{8, 3, "main"};
  auto EntryTwo = LineColFun{8, 8, "main"};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

TEST_F(IFDSTaintAnalysisTest, TaintTest_LibSummary_01) {
  initialize({PathToLlFiles + "dummy_source_sink/taint_lib_sum_01_cpp_dbg.ll"});
  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
  TaintSolver.solve();

  auto Entry = LineColFun{8, 3, "main"};
  auto EntryTwo = LineColFunOp{8, 8, "main", llvm::Instruction::FPToSI};
  GroundTruthTy GroundTruth{{Entry, {EntryTwo}}};

  compare(TaintProblem->Leaks, GroundTruth);
}

int main(int Argc, char **Argv) {
  ::testing::InitGoogleTest(&Argc, Argv);
  return RUN_ALL_TESTS();
}
